"""
简化版模型推理接口
专注于高效的实时推理
"""
import os
import logging
from typing import Dict, List, Optional, Union, Tuple
import torch
import numpy as np


class SimpleMahjongInference:
    """
    简化版麻将推理接口类
    """
    def __init__(self, 
                 model_path: str,
                 device: Optional[str] = None):
        """
        初始化简化版推理接口
        
        Args:
            model_path: 模型路径
            device: 运行设备
        """
        # 设置日志
        self.logger = self._setup_logging()
        self.logger.info(f"初始化简化版推理接口，模型路径: {model_path}")
        
        # 设备
        if device is None:
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.device = device
        self.logger.info(f"使用设备: {self.device}")
        
        # 加载模型
        self.model = self._load_model(model_path)
        
        # 动作映射
        self.action_map = self._setup_action_map()
    
    def _setup_logging(self) -> logging.Logger:
        """
        设置日志
        
        Returns:
            日志记录器
        """
        logger = logging.getLogger('SimpleMahjongInference')
        logger.setLevel(logging.INFO)
        
        # 控制台日志
        if not logger.handlers:
            console_handler = logging.StreamHandler()
            console_handler.setLevel(logging.INFO)
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            console_handler.setFormatter(formatter)
            logger.addHandler(console_handler)
        
        return logger
    
    def _load_model(self, model_path: str) -> torch.nn.Module:
        """
        加载模型
        
        Args:
            model_path: 模型路径
            
        Returns:
            模型实例
        """
        self.logger.info(f"加载模型: {model_path}")
        
        try:
            # 直接加载整个模型
            model = torch.jit.load(model_path, map_location=self.device)
            
            # 设置为评估模式
            model.eval()
            
            self.logger.info("模型加载成功")
            return model
        except Exception as e:
            self.logger.error(f"模型加载失败: {str(e)}")
            raise
    
    def _setup_action_map(self) -> List[str]:
        """
        设置动作映射
        
        Returns:
            动作映射列表
        """
        return [
            'pass',      # 过牌
            'play',      # 打牌
            'chi',       # 吃
            'peng',      # 碰
            'gang',      # 杠
            'hu'         # 和
        ]
    
    def preprocess_input(self, hand_tiles: List[int], 
                         visible_tiles: List[int],
                         last_tile: Optional[int] = None) -> torch.Tensor:
        """
        预处理输入数据
        
        Args:
            hand_tiles: 手牌
            visible_tiles: 可见的牌（其他玩家打出的牌）
            last_tile: 最后一张牌
            
        Returns:
            特征张量
        """
        # 创建手牌特征（简化版）
        hand_feature = np.zeros(108, dtype=np.float32)  # 四川麻将最多108张牌
        for tile in hand_tiles:
            if 0 <= tile < 108:
                hand_feature[tile] += 1
        
        # 创建可见牌特征
        visible_feature = np.zeros(108, dtype=np.float32)
        for tile in visible_tiles:
            if 0 <= tile < 108:
                visible_feature[tile] += 1
        
        # 创建最后一张牌特征
        last_tile_feature = np.zeros(108, dtype=np.float32)
        if last_tile is not None and 0 <= last_tile < 108:
            last_tile_feature[last_tile] = 1
        
        # 合并特征
        features = np.concatenate([hand_feature, visible_feature, last_tile_feature])
        
        # 转换为张量
        tensor = torch.tensor(features).to(self.device)
        tensor = tensor.unsqueeze(0)  # 添加批次维度
        
        return tensor
    
    def predict_action(self, hand_tiles: List[int], 
                      visible_tiles: List[int],
                      last_tile: Optional[int] = None) -> str:
        """
        预测动作
        
        Args:
            hand_tiles: 手牌
            visible_tiles: 可见的牌
            last_tile: 最后一张牌
            
        Returns:
            预测的动作
        """
        # 预处理输入
        features = self.preprocess_input(hand_tiles, visible_tiles, last_tile)
        
        # 预测
        with torch.no_grad():
            output = self.model(features)
            _, predicted = torch.max(output, 1)
            action_idx = predicted.item()
        
        # 返回动作名称
        if 0 <= action_idx < len(self.action_map):
            return self.action_map[action_idx]
        else:
            return 'pass'  # 默认返回过牌
    
    def predict_with_probs(self, hand_tiles: List[int], 
                          visible_tiles: List[int],
                          last_tile: Optional[int] = None) -> Tuple[str, np.ndarray]:
        """
        预测动作并返回概率
        
        Args:
            hand_tiles: 手牌
            visible_tiles: 可见的牌
            last_tile: 最后一张牌
            
        Returns:
            (预测的动作, 概率分布)
        """
        # 预处理输入
        features = self.preprocess_input(hand_tiles, visible_tiles, last_tile)
        
        # 预测
        with torch.no_grad():
            output = self.model(features)
            
            # 计算概率
            probs = torch.nn.functional.softmax(output, dim=1)
            probs_array = probs.cpu().numpy()[0]
            
            # 获取预测动作
            _, predicted = torch.max(output, 1)
            action_idx = predicted.item()
        
        # 返回动作名称和概率
        action_name = self.action_map[action_idx] if 0 <= action_idx < len(self.action_map) else 'pass'
        return action_name, probs_array


class OnlineMahjongPlayer:
    """
    在线麻将玩家类
    简化的接口，专为在线游戏设计
    """
    def __init__(self, model_path: str, 
                 player_id: int = 0,
                 device: Optional[str] = None):
        """
        初始化在线玩家
        
        Args:
            model_path: 模型路径
            player_id: 玩家ID
            device: 运行设备
        """
        # 加载推理接口
        self.inference = SimpleMahjongInference(model_path=model_path, device=device)
        self.player_id = player_id
        self.logger = self.inference.logger
        
        # 状态跟踪
        self.hand_tiles = []
        self.visible_tiles = []
        self.last_tile = None
        self.round = 0
    
    def update_state(self, hand_tiles: List[int], 
                    visible_tiles: List[int],
                    last_tile: Optional[int] = None) -> None:
        """
        更新玩家状态
        
        Args:
            hand_tiles: 手牌
            visible_tiles: 可见的牌
            last_tile: 最后一张牌
        """
        self.hand_tiles = hand_tiles
        self.visible_tiles = visible_tiles
        self.last_tile = last_tile
        self.round += 1
        
        self.logger.debug(f"玩家 {self.player_id} 状态更新，回合: {self.round}")
    
    def make_move(self) -> str:
        """
        做出移动
        
        Returns:
            要执行的动作
        """
        return self.inference.predict_action(
            self.hand_tiles,
            self.visible_tiles,
            self.last_tile
        )
    
    def make_move_with_info(self) -> Dict:
        """
        做出移动并返回详细信息
        
        Returns:
            包含动作和概率的字典
        """
        # 预测
        action, probs = self.inference.predict_with_probs(
            self.hand_tiles,
            self.visible_tiles,
            self.last_tile
        )
        
        # 构建结果
        result = {
            'player_id': self.player_id,
            'action': action,
            'action_prob': float(probs[self.inference.action_map.index(action)]),
            'all_actions_probs': {
                self.inference.action_map[i]: float(probs[i]) 
                for i in range(len(self.inference.action_map))
            },
            'hand_size': len(self.hand_tiles),
            'round': self.round
        }
        
        return result
    
    def should_play_tile(self, available_tiles: List[int] = None) -> int:
        """
        决定要打出的牌
        
        Args:
            available_tiles: 可用的牌（如果不提供则使用手牌）
            
        Returns:
            要打出的牌
        """
        # 这里可以添加简单的规则或使用另一个模型
        # 简化版本返回第一张牌
        if not available_tiles:
            available_tiles = self.hand_tiles
        
        return available_tiles[0] if available_tiles else -1


def export_model_for_inference(model, output_path: str) -> None:
    """
    导出模型用于推理
    
    Args:
        model: 模型实例
        output_path: 输出路径
    """
    # 创建目录
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    # 导出为TorchScript
    scripted_model = torch.jit.script(model)
    scripted_model.save(output_path)
    
    print(f"模型已导出到: {output_path}")


def main() -> None:
    """
    主函数（用于测试）
    """
    import argparse
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='简化版麻将推理接口')
    parser.add_argument('--model', type=str, help='模型路径')
    parser.add_argument('--device', type=str, help='运行设备')
    
    args = parser.parse_args()
    
    if not args.model:
        print("请提供模型路径")
        return
    
    # 创建推理接口
    inference = SimpleMahjongInference(
        model_path=args.model,
        device=args.device
    )
    
    # 测试数据
    test_hand = [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5, 6, 7]
    test_visible = [8, 9, 10]
    test_last = 8
    
    # 预测
    action = inference.predict_action(test_hand, test_visible, test_last)
    print(f"预测动作: {action}")
    
    # 预测并返回概率
    action, probs = inference.predict_with_probs(test_hand, test_visible, test_last)
    print(f"\n预测动作: {action}")
    print("动作概率:")
    for i, prob in enumerate(probs):
        print(f"  {inference.action_map[i]}: {prob:.4f}")
    
    # 创建在线玩家
    player = OnlineMahjongPlayer(
        model_path=args.model,
        player_id=0,
        device=args.device
    )
    
    # 更新状态
    player.update_state(test_hand, test_visible, test_last)
    
    # 做出移动
    move = player.make_move()
    print(f"\n玩家移动: {move}")
    
    # 做出移动并获取详细信息
    move_info = player.make_move_with_info()
    print("\n详细移动信息:")
    print(f"玩家ID: {move_info['player_id']}")
    print(f"动作: {move_info['action']}")
    print(f"动作概率: {move_info['action_prob']:.4f}")


if __name__ == '__main__':
    main()
